package mysql

/*
 * mysql wrapper.
 */
import (
	"database/sql"
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

// default pool parameter.
const (
	gMaxConnSize     = 32
	gMaxIdleConnSize = 16
	gMaxLifeTime     = 2 * time.Hour
)

// db config.
type DBConf struct {
	Host           string
	Port           uint16
	Database       string
	User           string
	Password       string
	Charset        string
	ConnectTimeout uint32
	ReadTimeout    uint32
	WriteTimeout   uint32
	HostList       []string
	MaxOpenConn    int
	MaxIdleConn    int
	MaxLifeTime    time.Duration
}

func NewDBConf() *DBConf {
	return &DBConf{
		Port:           3306,
		ConnectTimeout: 5,
		MaxOpenConn:    gMaxConnSize,
		MaxIdleConn:    gMaxIdleConnSize,
		MaxLifeTime:    gMaxLifeTime,
	}
}

func (conf *DBConf) Set(m map[string]string) {
	if v, ok := m["dbhost"]; ok {
		conf.Host = v
	}

	if v, ok := m["dbport"]; ok {
		if port, err := strconv.Atoi(v); err == nil {
			conf.Port = uint16(port)
		}
	}
	if v, ok := m["dbname"]; ok {
		conf.Database = v
	}
	if v, ok := m["dbuser"]; ok {
		conf.User = v
	}
	if v, ok := m["dbpass"]; ok {
		conf.Password = v
	}
	if v, ok := m["charset"]; ok {
		conf.Charset = v
	}
	if v, ok := m["conntimeout"]; ok {
		if t, err := strconv.Atoi(v); err == nil {
			conf.ConnectTimeout = uint32(t)
		}
	}
	if v, ok := m["readtimeout"]; ok {
		if t, err := strconv.Atoi(v); err == nil {
			conf.ReadTimeout = uint32(t)
		}
	}
	if v, ok := m["writetimeout"]; ok {
		if t, err := strconv.Atoi(v); err == nil {
			conf.WriteTimeout = uint32(t)
		}
	}
	if v, ok := m["hostlist"]; ok {
		conf.HostList = strings.Split(v, ",;")
		for idx, host := range conf.HostList {
			conf.HostList[idx] = strings.TrimSpace(host)
		}
	}
}

// mysql executor proxy.
type MysqlExecutorProxy interface {
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryEx(query string, args ...interface{}) (*MysqlResult, error)
	Exec(query string, args ...interface{}) (sql.Result, error)
	Commit() error
	Rollback() error
}

// mysql client with tx.
type MysqlTxClient struct {
	*sql.Tx
}

func (mtc *MysqlTxClient) Query(query string, args ...interface{}) (*sql.Rows, error) {
	return mtc.Tx.Query(query, args...)
}
func (mtc *MysqlTxClient) QueryEx(query string, args ...interface{}) (*MysqlResult, error) {
	if rows, err := mtc.Query(query, args...); err != nil {
		return nil, err
	} else {
		return NewMysqlResult(rows), nil
	}
}
func (mtc *MysqlTxClient) Exec(query string, args ...interface{}) (sql.Result, error) {
	return mtc.Tx.Exec(query, args...)
}
func (mtc *MysqlTxClient) Commit() error {
	return mtc.Tx.Commit()
}
func (mtc *MysqlTxClient) Rollback() error {
	return mtc.Tx.Rollback()
}

// mysql client.
type MysqlClient struct {
	conf    *DBConf
	db      *sql.DB
	hostIdx int
}

func (cli *MysqlClient) getConnString() (connstr string, addr string) {
	if cli.hostIdx == 0 {
		if cli.conf.Host != "" {
			addr = fmt.Sprintf("%s:%d", cli.conf.Host, cli.conf.Port)
		}
	} else if cli.hostIdx <= len(cli.conf.HostList) {
		addr = cli.conf.HostList[cli.hostIdx-1]
	}
	if addr == "" {
		return
	}

	var builder strings.Builder
	fmt.Fprintf(&builder, "%s:%s@tcp(%s)/%s?timeout=%ds", cli.conf.User, cli.conf.Password, addr, cli.conf.Database, cli.conf.ConnectTimeout)
	if cli.conf.ReadTimeout != 0 {
		fmt.Fprintf(&builder, "&readTimeout=%ds", cli.conf.ReadTimeout)
	}
	if cli.conf.WriteTimeout != 0 {
		fmt.Fprintf(&builder, "&writeTimeout=%ds", cli.conf.WriteTimeout)
	}
	if cli.conf.Charset != "" {
		fmt.Fprintf(&builder, "&charset=%s", cli.conf.Charset)
	}
	connstr = builder.String()
	return
}

func (cli *MysqlClient) connect(connstr string) (*sql.DB, error) {
	db, err := sql.Open("mysql", connstr)
	if err != nil {
		return nil, err
	}

	// set default value if possible.
	if cli.conf.MaxOpenConn == 0 {
		cli.conf.MaxOpenConn = gMaxConnSize
	}
	if cli.conf.MaxIdleConn == 0 {
		cli.conf.MaxIdleConn = gMaxIdleConnSize
	}
	if cli.conf.MaxLifeTime == 0 {
		cli.conf.MaxLifeTime = gMaxLifeTime
	}
	db.SetMaxOpenConns(cli.conf.MaxOpenConn)
	db.SetMaxIdleConns(cli.conf.MaxIdleConn)
	db.SetConnMaxLifetime(cli.conf.MaxLifeTime)

	err = db.Ping()
	if err != nil {
		db.Close()
		return nil, err
	}
	return db, err
}

// 支持cluster连接方式
func (cli *MysqlClient) ConnectCluster() error {
	if cli.db != nil {
		return nil
	}

	// 构建MySQL集群的DSN
	user := cli.conf.User
	password := cli.conf.Password
	database := cli.conf.Database
	timeOut := cli.conf.ConnectTimeout
	clusterAddr := strings.Join(cli.conf.HostList, ",")
	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?timeout=%ds", user, password, clusterAddr, database, timeOut)
	if cli.conf.ReadTimeout != 0 {
		dsn = fmt.Sprintf("%s&readTimeout=%ds", dsn, cli.conf.ReadTimeout)
	}
	if cli.conf.WriteTimeout != 0 {
		dsn = fmt.Sprintf("%s&writeTimeout=%ds", dsn, cli.conf.WriteTimeout)
	}
	if cli.conf.Charset != "" {
		dsn = fmt.Sprintf("%s&charset=%s", dsn, cli.conf.Charset)
	}

	var err error
	cli.db, err = cli.connect(dsn)
	return err
}

// 普通的连接方式
func (cli *MysqlClient) Connect() error {
	if cli.db != nil {
		return nil
	}

	var errmsg []string
	for try := len(cli.conf.HostList) + 1; try > 0; try-- {
		if cli.hostIdx > len(cli.conf.HostList) {
			cli.hostIdx = 0
		}
		connstr, addr := cli.getConnString()
		if connstr == "" {
			cli.hostIdx++
			continue
		}

		db, err := cli.connect(connstr)
		if err != nil {
			errmsg = append(errmsg, fmt.Sprintf("connect %s failed: %s", addr, err.Error()))
			cli.hostIdx++
			continue
		}

		cli.db = db
		return nil
	}
	return fmt.Errorf("can not find valid mysql endpoint: %s", strings.Join(errmsg, "; "))
}

func (cli *MysqlClient) Close() {
	if cli.db == nil {
		return
	}
	cli.db.Close()
	cli.db = nil
}

// query interface.
func (cli *MysqlClient) Query(query string, args ...interface{}) (*sql.Rows, error) {
	return cli.db.Query(query, args...)
}
func (cli *MysqlClient) QueryEx(query string, args ...interface{}) (*MysqlResult, error) {
	rows, err := cli.Query(query, args...)
	if err != nil {
		return nil, err
	}
	return NewMysqlResult(rows), nil
}

// execute interface.
func (cli *MysqlClient) Exec(query string, args ...interface{}) (sql.Result, error) {
	return cli.db.Exec(query, args...)
}

func (cli *MysqlClient) Commit() error {
	return nil
}
func (cli *MysqlClient) Rollback() error {
	return nil
}

// mysql client with tx.
func (cli *MysqlClient) CreateTx() (MysqlExecutorProxy, error) {
	if tx, err := cli.db.Begin(); err != nil {
		return nil, err
	} else {
		return &MysqlTxClient{
			Tx: tx,
		}, nil
	}
}
func (cli *MysqlClient) Prepare(query string) (*sql.Stmt, error) {
	return cli.db.Prepare(query)
}

func (cli *MysqlClient) GetHandle() *sql.DB {
	return cli.db
}

func EscapeString(v string) []byte {
	pos := 0
	buf := make([]byte, len(v)*2)

	for i := 0; i < len(v); i++ {
		c := v[i]
		switch c {
		case 0:
			buf[pos] = '\\'
			buf[pos+1] = '0'
			pos += 2
		case '\n':
			buf[pos] = '\\'
			buf[pos+1] = 'n'
			pos += 2
		case '\r':
			buf[pos] = '\\'
			buf[pos+1] = 'r'
			pos += 2
		case 0x1a:
			buf[pos] = '\\'
			buf[pos+1] = 'Z'
			pos += 2
		case '\'':
			buf[pos] = '\\'
			buf[pos+1] = '\''
			pos += 2
		case '"':
			buf[pos] = '\\'
			buf[pos+1] = '"'
			pos += 2
		case '\\':
			buf[pos] = '\\'
			buf[pos+1] = '\\'
			pos += 2
		default:
			buf[pos] = c
			pos++
		}
	}
	return buf[:pos]
}

func EscapeBytes(v []byte) []byte {
	pos := 0
	buf := make([]byte, len(v)*2)

	for i := 0; i < len(v); i++ {
		c := v[i]
		switch c {
		case 0:
			buf[pos] = '\\'
			buf[pos+1] = '0'
			pos += 2
		case '\n':
			buf[pos] = '\\'
			buf[pos+1] = 'n'
			pos += 2
		case '\r':
			buf[pos] = '\\'
			buf[pos+1] = 'r'
			pos += 2
		case 0x1a:
			buf[pos] = '\\'
			buf[pos+1] = 'Z'
			pos += 2
		case '\'':
			buf[pos] = '\\'
			buf[pos+1] = '\''
			pos += 2
		case '"':
			buf[pos] = '\\'
			buf[pos+1] = '"'
			pos += 2
		case '\\':
			buf[pos] = '\\'
			buf[pos+1] = '\\'
			pos += 2
		default:
			buf[pos] = c
			pos++
		}
	}
	return buf[:pos]
}

// outside interface.
func NewMysqlClient(conf *DBConf) *MysqlClient {
	cli := &MysqlClient{
		conf: conf,
	}

	if len(conf.HostList) > 0 {
		randomInt := rand.Intn(len(conf.HostList))
		cli.hostIdx = randomInt
	}
	return cli
}

func init() {
	// random.
	rand.Seed(time.Now().UnixNano())
}
